#pragma once

#include <cmath>
#include <cstdint>

struct TermLocalStat{
    int title_cnt = 0;
    int content_cnt = 0;
};

struct DocStat{
    std::uint32_t length = 0;   //文档分词总数
};

struct CorpusStat{
    std::uint64_t N = 0;        // 文档总数
    std::uint64_t df = 0;       // 出现该词的文档数
    double avg_doc_len = 0.0;   // 平均文档长度
};

// Interface
class IWeightStrategy{
public:
    virtual ~IWeightStrategy() = default;
    virtual double compute(const TermLocalStat &term,
                           const DocStat &doc,
                           const CorpusStat &corpus) const = 0;
};

class TFIDFStrategy : public IWeightStrategy{
private:
    double _title_w = 0.0;
    double _content_w = 0.0;
    
public:
    TFIDFStrategy(double title_w = 10.0, double content_w = 1.0)
    : _title_w(title_w), _content_w(content_w)
    {}
    double compute(const TermLocalStat &term,
                   const DocStat &doc,//can be null
                   const CorpusStat &corpus) const override {
        //1.TF:字段加权 + log压缩
        double tf_raw = term.title_cnt * _title_w + term.content_cnt * _content_w;
        if(tf_raw <= 0) return 0.0;
        double tf = 1.0 + std::log(tf_raw + 1.0);
        //2.IDF
        double df = corpus.df;
        double N = corpus.N;
        double idf = std::log((N + 1.0) / (df + 1.0)) + 1.0;

        return tf * idf;
    }
};

class BM25Strategy : public IWeightStrategy {
private:
    double title_w_;
    double content_w_;
    double k1_;
    double b_;

public:
    BM25Strategy(double title_w = 10.0, double content_w = 1.0,
                 double k1 = 1.5, double b = 0.75)
        : title_w_(title_w), content_w_(content_w),
          k1_(k1), b_(b) {}

    double compute(const TermLocalStat& term,
                   const DocStat& doc,
                   const CorpusStat& corpus) const override {
        // 0. 基础数据安全检查
        if (corpus.N == 0 || corpus.df == 0) return 0.0;
        if (doc.length == 0 || corpus.avg_doc_len <= 0.0) return 0.0;

        // 1. 先做字段加权的 tf_raw
        int tf_raw = term.title_cnt * title_w_ + term.content_cnt * content_w_;
        if (tf_raw <= 0) return 0.0;
        double tf = static_cast<double>(tf_raw);

        // 2. BM25 的 idf
        double N = static_cast<double>(corpus.N);
        double df = static_cast<double>(corpus.df);
        double idf = std::log((N - df + 0.5) / (df + 0.5) + 1.0);

        // 3. 文档长度归一化
        double dl  = static_cast<double>(doc.length);
        double avdl = corpus.avg_doc_len;

        double denom = tf + k1_ * (1.0 - b_ + b_ * dl / avdl);
        double score = idf * (tf * (k1_ + 1.0)) / denom;
        return score;
    }
};